import pandas as pd
import numpy as np
import math

DST = '11.11.11.2'

def load_tcpdump_pcap_csv(filename, length = None):
    if length == None:
        df_data = pd.read_csv(filename)
    else:
        df_data = pd.read_csv(filename)[0:length]
    return df_data


def get_packets_in_flight(filename, length = None):
    df_data = load_tcpdump_pcap_csv(filename, length)
    inflight_packets = 0
    inflight_packets_list = []
    current_time_list = []
    last_len = 0
    last_ack = 0

    for index, row in df_data.iterrows():
        time_relative = row['frame.time_relative']
        current_time_list.append(time_relative)
        ack = row['tcp.ack']
        tcp_len = row['tcp.len']
    
        if ack == 1:
            inflight_packets += tcp_len
        else:
            if ack - last_ack >= 0:
                inflight_packets -= ack - last_ack
                last_ack = ack
            else:
                print('tmp debug', ack - last_ack, ack, last_ack)
        inflight_packets_list.append(inflight_packets)
    
    return current_time_list, inflight_packets_list


def get_delay_list(filename, length = None):
    df_data = load_tcpdump_pcap_csv(filename, length)
    delay_list = []
    current_time_list = []

    record_need_ack_dict = {}

    last_ack = 0

    for index, row in df_data.iterrows():
        time_relative = row['frame.time_relative']
        seq = row['tcp.seq']
        ack = row['tcp.ack']
        tcp_len = row['tcp.len']
    
        if ack == 1:
            if row['ip.dst'] == DST:
                need_ack = seq + tcp_len
                record_need_ack_dict[need_ack] = time_relative
        else:
            if row['ip.dst'] != DST:
                if ack - last_ack >= 0:
                    for key in list(record_need_ack_dict.keys()):
                        if key < ack:
                            del record_need_ack_dict[key]
                        elif key == ack:
                            delay_list.append(time_relative - record_need_ack_dict[key])
                            current_time_list.append(time_relative)
                            del record_need_ack_dict[key]

                    last_ack = ack
                else:
                    print('tmp debug', ack - last_ack, ack, last_ack)
        
    return current_time_list, delay_list


def get_packets_in_flight_old2(filename, length = None, interval=0.01):
    df_data = load_tcpdump_pcap_csv(filename, length)
    inflight_packets = 0
    inflight_packets_list = []
    current_time_list = []
    record_need_ack_dict = {}

    interval_cnt = interval
    for index, row in df_data.iterrows():
        if row['ip.dst'] == DST:    # mm -> hombre
            seq = row['tcp.seq']
            time_relative = row['frame.time_relative']
            frame_len = row['frame.len']
            record_need_ack_dict[seq] = frame_len
            inflight_packets += frame_len

            if time_relative >= interval_cnt:
                inflight_packets_list.append(inflight_packets)
                current_time_list.append(time_relative)
                interval_cnt += interval
        else:    # hombre -> mm
            ack = row['tcp.ack']
            for key in list(record_need_ack_dict.keys()):
                if record_need_ack_dict[key] <= ack:
                    inflight_packets -= record_need_ack_dict[key]
                    del record_need_ack_dict[key]
            time_relative = row['frame.time_relative']
            if time_relative >= interval_cnt:
                current_time_list.append(time_relative)
                inflight_packets_list.append(inflight_packets)
                interval_cnt += interval
    return current_time_list, inflight_packets_list


def get_packets_in_flight_old(filename, length = None):
    df_data = load_tcpdump_pcap_csv(filename, length)
    inflight_packets = 0
    inflight_packets_list = []
    current_time_list = []

    record_ack_list = []
    last_time = 0
    for index, row in df_data.iterrows():
        if row['ip.dst'] == DST:    # mm -> hombre
            # Only caculate ACK flag and ACK-PSH flag 
            if row['tcp.flags'] == '0x00000010' or row['tcp.flags'] == '0x00000018':
                seq = row['tcp.seq']
                time_relative = row['frame.time_relative']

                if len(record_ack_list) > 0:
                    while time_relative > record_ack_list[0][0] and record_ack_list[0][0] > last_time:
                        inflight_packets -= record_ack_list[0][1]
                        inflight_packets_list.append(inflight_packets)
                        current_time_list.append(record_ack_list[0][0])
                        if len(record_ack_list) > 0:
                            del record_ack_list[0]
                        if len(record_ack_list) == 0:
                            break

                potential_ack_df = df_data[
                    (df_data['tcp.ack'] == seq)
                    & (df_data['ip.dst'] != DST)
                    & (df_data['tcp.flags'] == '0x00000010')
                    & (df_data['frame.time_relative'] > time_relative)
                ]
                # Filter every packets without ack. (usually every first [P.] packet.)
                if not potential_ack_df.empty:
                    snd_size = row['frame.len']
                    inflight_packets += snd_size
                    inflight_packets_list.append(inflight_packets)
                    current_time_list.append(time_relative)

                    # Only select latest ack packet.
                    potential_ack = potential_ack_df.iloc[0]
                    potential_ack_time = potential_ack['frame.time_relative']
                    record_ack_list.append((potential_ack_time, snd_size))
                    record_ack_list = sorted(record_ack_list, key = lambda x: x[0])
                last_time = time_relative
    return current_time_list, inflight_packets_list


def get_delay_list_old(filename, length = None):
    df_data = load_tcpdump_pcap_csv(filename, length)
    delay_list = []
    current_time_list = []
    for index, row in df_data.iterrows():
        if row['ip.dst'] == DST:    # mm -> hombre
            # Only caculate ACK flag and ACK-PSH flag 
            if row['tcp.flags'] == '0x00000010' or row['tcp.flags'] == '0x00000018':
                seq = row['tcp.seq']
                if seq > 500:
                    time_relative = row['frame.time_relative']
                    potential_ack_df = df_data[
                        (df_data['tcp.ack'] == seq)
                        & (df_data['ip.dst'] != DST)
                        & (df_data['tcp.flags'] == '0x00000010')
                        & (df_data['frame.time_relative'] > time_relative)
                    ]
                    # Filter every packets without ack. (usually every first [P.] packet.)
                    if not potential_ack_df.empty:
                        # Only select latest ack packet.
                        potential_ack = potential_ack_df.iloc[0]
                        potential_ack_time = potential_ack['frame.time_relative']
                        delay_list.append(potential_ack_time - time_relative)
                        current_time_list.append(time_relative)
    return current_time_list, delay_list


def get_throughput_list(filename, length = None, interval = 0.1):
    df_data = load_tcpdump_pcap_csv(filename, length)
    snd_data = df_data[df_data['ip.dst'] == DST]
    max_time = math.ceil(snd_data.iloc[-1]['frame.time_relative'])
    min_time = math.ceil(snd_data.iloc[0]['frame.time_relative'])

    time_list = []
    throughput_list = []

    for i in np.arange(min_time, max_time, interval):
        throughput = snd_data[(snd_data['frame.time_relative'] > i - 1) & (snd_data['frame.time_relative'] < i)]['frame.len'].sum()
        time_list.append(i)
        throughput_list.append(throughput)
    return time_list, throughput_list


def get_loss_list_old(filename, expected_max_rtt = 0.1):
    df_data = load_tcpdump_pcap_csv(filename)
    snd_data = df_data[df_data['ip.dst'] == DST]
    rcv_data = df_data[df_data['ip.dst'] != DST]

    loss_time_list = []

    for index, row in snd_data.iterrows():
        seq = row['tcp.seq']
        tcp_len = row['tcp.len']
        relative_time = row['frame.time_relative']

        if seq > 1:
            need_ack = seq + tcp_len
            ack_df = rcv_data[(rcv_data['tcp.ack'] == need_ack) & (rcv_data['frame.time_relative'] > relative_time)]
            if len(ack_df) > 0: 
                time_diff = ack_df.iloc[0]['frame.time_relative'] - relative_time
                if time_diff > expected_max_rtt:
                    # packet loss
                    loss_time_list.append(relative_time)
            else:
                # packet loss
                loss_time_list.append(relative_time)

    return loss_time_list


def get_loss_list(filename, expected_max_rtt = 0):
    df_data = load_tcpdump_pcap_csv(filename)
    snd_data = df_data[df_data['ip.dst'] == DST]

    loss_time_list = []

    for index, row in snd_data.iterrows():
        seq = row['tcp.seq']
        tcp_len = row['tcp.len']
        relative_time = row['frame.time_relative']

        if seq > 1:
            # repeat_df = snd_data[(snd_data['tcp.seq'] == seq) & (snd_data['tcp.len'] == tcp_len)]
            repeat_df = snd_data[(snd_data['tcp.seq'] == seq)]
            if len(repeat_df) > 1:
                for _, repeat_row in repeat_df.iterrows():
                    time_diff = repeat_row['frame.time_relative'] - relative_time
                    if time_diff > expected_max_rtt:
                        # packet loss
                        loss_time_list.append(relative_time)

    return loss_time_list


def get_loss_list_from_two_tcpdump(filename_a, filename_b):
    df_data_a = load_tcpdump_pcap_csv(filename_a)
    snd_data_a = df_data_a[df_data_a['ip.dst'] == DST]
    df_data_b = load_tcpdump_pcap_csv(filename_b)
    snd_data_b = df_data_b[df_data_b['ip.dst'] == DST]

    loss_time_list = []
    for index, row in snd_data_a.iterrows():
        seq = row['tcp.seq']
        relative_time = row['frame.time_relative']

        if seq > 1:
            check_b = snd_data_b[snd_data_b['tcp.seq'] == seq]
            if len(check_b) < 1:
                loss_time_list.append(relative_time)
    return loss_time_list

if __name__ == '__main__':
    # current_time_list, inflight_packets_list = get_packets_in_flight('cubic_md5_100Mb_lat10ms_uplinkloss1%.pcap.csv')
    # current_time_list, delay_list = get_delay_list('cubic_md5_100Mb_lat10ms_uplinkloss1%.pcap.csv')
    # current_time_list, throughput_list = get_throughput_list('cubic_md5_100Mb_lat10ms_uplinkloss1%.pcap.csv', None, 0.1)
    # print(current_time_list, throughput_list)
    # loss_time_list = get_loss_list('replay_wireless_cubic.pcap.csv', 0)
    # loss_time_list = get_loss_list('loss_edith.pcap.csv', 0.001)
    # loss_time_list = get_loss_list_from_two_tcpdump('loss_ga_eth0.pcap.csv', 'loss_ga_eth1.pcap.csv')
    loss_time_list = get_loss_list_from_two_tcpdump('loss_Ga_eth0_mm1_cubic.pcap.csv', 'loss_Ga_eth1_mm1_cubic.pcap.csv')
    print(loss_time_list)
    print(len(loss_time_list))


